# -*- coding: utf-8 -*-

# 导入pyspark
from pyspark import SparkContext

# 导入决策树回归
from pyspark.mllib.tree import DecisionTree, DecisionTreeModel
from pyspark.mllib.regression import LabeledPoint
import numpy as np

# 导入可视化
import matplotlib
matplotlib.use('Agg') # 不回显
import matplotlib.pyplot as plt
# import matplotlib.pylab as plb

# 导入数据库操作
import MySQLdb

# 函数-决策树模型创建特征向量
def extract_features_dt(record):
	return np.array(map(float, record[2:14]))

# 函数-提取实际租借车数
def extract_label(record):
	return float(record[-1])

# 函数-平方误差
def squared_error(actual, pred):
	return (pred - actual)**2

# 函数-绝对误差
def abs_error(actual, pred):
	return np.abs(pred - actual)

# 函数-平方对数误差
def squared_log_error(pred, actual):
	return (np.log(pred + 1) - np.log(actual + 1))**2

# 函数-画出比较
def evaluate_dt_draw(model, test):

	preds = model.predict(test.map(lambda p: p.features))
	actual = test.map(lambda p: p.label)
	# 画图
	plt.figure(num=1, figsize=(8,6))
	plt.title('Decision Tree Regression Model Test', size=14)
	plt.xlabel('Actuals', size=14)
	plt.ylabel('Predicts', size=14)
	actual_take = actual.take(2)
	preds_take = preds.take(2)
	plt.plot(actual_take, actual_take)
	plt.plot(actual_take, preds_take, "r+")
	# text 点
	for i in range(0, 2):
		arrShow = (int(actual_take[i]), round(preds_take[i], 2))
		textShow = str(arrShow)
		plt.text(actual_take[i], preds_take[i], textShow, color="red", fontsize=6)

	plt.savefig('DecisionTreeTest.png', format='png')
	# 返回 rsmle
	tp = actual.zip(preds)
	mse = tp.map(lambda (t, p): squared_error(t, p)).mean()
	rmsle = np.sqrt(tp.map(lambda (t, p): squared_log_error(t,p)).mean())
	sql = "insert into queue(name, status, mse, rmsle) values('%s', %d, %f, %f)"%("DecisionTreeRegression", 1, mse, rmsle)
	insertNewRecordForDecisionTreeRegression(sql)
	return rmsle

# 函数-操作数据库
def insertNewRecordForDecisionTreeRegression(sql):
	# 连接数据库
	try:
		conn = MySQLdb.connect(host='localhost',user='root',passwd='root',db='sparkDemo')
	except Exception, e:
		print e
	# cursor对象操作数据库
	cursor = conn.cursor()
	# 执行SQL
	try:
		cursor.execute(sql)
	except Exception, e:
		print e
	# 提交操作
	conn.commit()
	# 关闭指针
	cursor.close()
	# 关闭数据库
	conn.close()
	return 1

# main函数部分
sc = SparkContext("yarn-client", "Decision Tree Regression Spark App")

path = "DecisionTreeModel"
sampleModel = DecisionTreeModel.load(sc, path)

# 加载数据集
path = "file:///var/www/html/web/regressionTest.csv" # 数据集位置
raw_data = sc.textFile(path) # 数据集
records = raw_data.map(lambda x: x.split(",")) # 通过逗号分隔
data_dt = records.map(lambda r: LabeledPoint(extract_label(r), extract_features_dt(r)))
evaluate_dt_draw(sampleModel, data_dt)

sc.stop()